import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
import os, sys
from PIL import Image

from diffusers import (
    AutoencoderKL,
    UNet2DConditionModel,
    EulerDiscreteScheduler
)
import random

from hydit.kolors_models.pipelines.pipeline_stable_diffusion_xl_chatglm_256_inpainting import StableDiffusionXLInpaintPipeline
from hydit.kolors_models.models.modeling_chatglm import ChatGLMModel
from hydit.kolors_models.models.tokenization_chatglm import ChatGLMTokenizer

def infer(image_path_list, mask_path_list, prompt_list, save_path):

    ckpt_dir = '/root/work/filestorage/liuxin/weight/Kolors-Inpainting'
    text_encoder = ChatGLMModel.from_pretrained(
        f'{ckpt_dir}/text_encoder',
        torch_dtype=torch.float16).half()
    tokenizer = ChatGLMTokenizer.from_pretrained(f'{ckpt_dir}/text_encoder')
    vae = AutoencoderKL.from_pretrained(f"{ckpt_dir}/vae", revision=None).half()
    scheduler = EulerDiscreteScheduler.from_pretrained(f"{ckpt_dir}/scheduler")

    unet = UNet2DConditionModel.from_pretrained(f"{ckpt_dir}/unet").half()

    pipe = StableDiffusionXLInpaintPipeline(
            vae=vae,
            text_encoder=text_encoder,
            tokenizer=tokenizer,
            unet=unet,
            scheduler=scheduler
    )
    
    pipe.to("cuda")
    pipe.enable_attention_slicing()

    for image_path, mask_path, prompt in zip(image_path_list, mask_path_list, prompt_list):
        image = Image.open(image_path).convert('RGB')
        mask_image = Image.open(mask_path).convert('RGB')
        width = mask_image.size[0] // 16 * 16
        height = mask_image.size[1] // 16 * 16
        generator = torch.Generator(device="cpu").manual_seed(1234)


        prompt = '生成一张浅黄色渐变背景图'
        result = pipe(
            prompt=prompt,
            image=image,
            mask_image = mask_image,
            height=1024,
            width=1024,
            negative_prompt="残缺的手指，畸形的手指，畸形的手，残肢，模糊，低质量",
            guidance_scale=6.0,
            num_inference_steps=50,
            generator=generator,
            strength = 1
        ).images[0]
        basename = 'result.jpg'
        result.save(f'{save_path}/{basename}')

if __name__ == '__main__':
    image_path_list = ['testimg/maskedimg.png']
    mask_path_list = ['testimg/maskimg.png']

    prompt_list = ['穿着美少女战士的衣服，一件类似于水手服风格的衣服，包括一个白色紧身上衣，前胸搭配一个大大的红色蝴蝶结。衣服的领子部分呈蓝色，并且有白色条纹。她还穿着一条蓝色百褶裙']

    # prompt_list = [''] * len(mask_path_list)
    save_path = 'testimg'
    os.makedirs(save_path, exist_ok=True)
    infer(image_path_list, mask_path_list, prompt_list, save_path)




